package com.lyl.myaiagent.optimize.filtering_retrieval;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.TranslationQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * @author lyl
 * @version 1.0
 * @CreateTime: 2025-08-19
 * 查询重写类和查询翻译类
 */
@Component
public class QueryRewriterAndTranslation {

    private final QueryTransformer rewriterqueryTransformer;

    private final QueryTransformer transqueryTransformer;

    public QueryRewriterAndTranslation(ChatModel dashscopeChatModel){
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        rewriterqueryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
        transqueryTransformer = TranslationQueryTransformer.builder()
                .chatClientBuilder(builder)
                .targetLanguage("zh")
                .build();
    }

    /**
     * 查询重写
     * @param prompt 提示词
     * @return 重写后的提示词
     */
    public String doWriter(String prompt){
        Query rewriter = rewriterqueryTransformer.apply(new Query(prompt));
        return rewriter.text();
    }

    /**
     * 查询翻译
     * @param prompt 提示词
     * @return 翻译后的提示词
     */
    public String doTrans(String prompt){
        Query reTrans = transqueryTransformer.apply(new Query(prompt));
        return reTrans.text();
    }
}
